# -*- mode: python -*-
###
#
# (c) 2000-2007 Wijnand Modderman <wijnand at tehmaze dot com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

__revision__ = '$Hg: ip.py,v 521b597a854c 2007/01/20 14:28:18 +0100 $'

# system imports
import re
import socket

# DDB imports
from ddb.tools import syntax
from ddb.threads import threaded
from ddb.register import register

# Twisted imports
from twisted.internet import protocol, reactor
from twisted.protocols import basic

# DDB exports
EXPORT={
    # module    flags needed
    #           global local
    'Ip':       (None, None)
    }

#
# TCP banner grabber
#
class BannerProtocol(basic.LineReceiver):
    delimiter = '\n'
    def lineReceived(self, line):
        self.factory.line = line.strip()
        self.transport.loseConnection()

class BannerFactory(protocol.ClientFactory):
    protocol = BannerProtocol
    failed   = "Connection reset before banner received"
    line     = failed

    def clientConnectionLost(self, connector, reason):
        self.manager.grabDone(self, self.line)

    def clientConnectionFailed(self, connector, reason):
        self.manager.grabErr(self, reason)

class BannerGrab(object):
    """I'm a very simple class that instigates a BannerFactory to grab the first
    line received on the socket and return immediately."""

    timeout = 5

    def __init__(self, event, callback, errback):
        self.event    = event
        self.callback = callback
        self.errback  = errback
        self.factory  = BannerFactory()
        self.factory.manager = self

    def grab(self, host, port):
        self.host = host
        self.port = port
        reactor.connectTCP(self.host, self.port, self.factory, timeout=self.timeout)

    def grabDone(self, factory, line):
        self.callback(self, line)

    def grabErr(self, factory, reason):
        self.errback(self, reason)

#
# The module
#

class Ip:
    """Internet Protocol functions."""

    _re_hexip = re.compile('^[\da-f]{8}$', re.IGNORECASE)

    def __init__(self):
        self.reg     = register.init(self)
        self.context = []
        self.timeout = self.reg.get('tcptimeout', *self.context)
        if not self.timeout:
            self.timeout = 5
            self.reg.set('tcptimeout', self.timeout, *self.context)
        else:
            self.timeout = self.timeout[0] # because the register returns lists always

    # Because we are going to do some blocking lookups here, fork this call into
    # a thread, so that we can continue to respond :-)
    @threaded
    def dns(self, event):
        """<host|ip>

        performs a DNS lookup an ip or hostname."""
        if not event.args:
            event.reply(syntax())
        else:
            event.reply('Attempting to look up %s' % event.args)
            is_a   = None
            result = None
            # If we support IPv6 ...
            if socket.has_ipv6:
                # ... then check if this is an IPv6 ip
                try:
                    socket.inet_pton(socket.AF_INET6, event.args)
                    is_a = 'ipv6'
                except socket.error:
                    pass
            # Ah not an IPv6 ip ...
            if not is_a:
                # ... maybe IPv4 ?
                try:
                    socket.inet_pton(socket.AF_INET, event.args)
                    is_a = 'ipv4'
                except socket.error:
                    pass
            # Not an ip, must be a hostname then
            if not is_a:
                is_a = 'host'
            # If it was an ip ...
            if is_a in ['ipv4', 'ipv6']:
                try:
                    # ... try to resolve it
                    result = socket.gethostbyaddr(event.args)
                    if result[1]:
                        result = 'primary: %s, aliases: %s' % \
                            (result[0], ', '.join(result[1]))
                    else:
                        result = result[0]
                    event.reply('%s ip %s resolves to: %s' % (is_a, event.args, result))
                except Exception, e:
                    try:
                        why = e[1]
                    except:
                        why = str(e)
                    event.reply('could not resolve %s address %s: %s' % \
                        (is_a, event.args, why))
            # Oh it's a host, lets resolve that
            elif is_a == 'host':
                try:
                    result = []
                    for info in socket.getaddrinfo(event.args, None):
                        if info[0] in [socket.AF_INET, socket.AF_INET6] and \
                            info[1] == socket.SOCK_STREAM:
                            ip = info[4][0]
                            if not ip in result:
                                result.append(ip)
                    if not result:
                        event.reply('could not resolve hostname %s: not found' % \
                            event.args)
                    else:
                        event.reply('%s resolves to: %s' % (event.args, ', '.join(result)))
                except Exception, e:
                    try:
                        why = e[1]
                    except:
                        why = str(e)
                    event.reply('could not resolve hostname %s: %s' % \
                        (event.args, why))
            else:
                event.reply('lookup failed, no valid data found')

    def hexip(self, event):
        """<ip|hex ip>

        returns the reverse of the given argument."""
        is_a = None
        if event.args.startswith('0x'):
            event.args = event.args[2:]
        if self._re_hexip.match(event.args):
            is_a = 'hexip'
        else:
            try:
                socket.inet_pton(socket.AF_INET, event.args)
                is_a = 'defip'
            except socket.error:
                pass
        if not is_a:
            event.reply(syntax())
            return
        if is_a == 'hexip':
            ip = []
            for i in range(4):
                ip.append(str(int(event.args[i*2:i*2+2], 16)))
            event.reply('%s = %s' % (event.args, '.'.join(ip)))
        else:
            test = event.args.split('.')
            ip = 16777216 * int(test[0]) + 65536 * int(test[1]) + 256 * int(test[2]) + int(test[3])
            event.reply('ip %s = %08x' % (event.args, ip))

    def ip(self, event):
        """depricated, use hexip."""
        event.reply(syntax())
        self.hexip(event)

    # Obviously blocking, we don't want that, so fork in thread
    @threaded
    def tcpbanner(self, event):
        """<[hostname]> <[port]>

        Attempt to grab the banner from tcp service on [hostname]:[port]."""
        # Attempt to parse the arguments
        try:
            host, port = event.args.split()
            port       = int(port)
        except:
            event.reply(syntax())
            return
        grabber = BannerGrab(event, self._tcpbanner, self._tcpbannererr)
        grabber.timeout = self.timeout
        grabber.grab(host, port)

    def _tcpbanner(self, grabber, line):
        grabber.event.reply('TCP banner on %s:%d: %s' % (grabber.host, grabber.port, line))

    def _tcpbannererr(self, grabber, reason):
        grabber.event.reply('TCP banner grabbing on %s:%s failed: %s' % (grabber.host, grabber.port, reason.getErrorMessage()))

# vim:ts=4:sw=4:et:
